﻿using System;

namespace exercise.树
{
    public class AVLTree<E> : BinarySearchTree<E>
    {
        protected Comparison<E> comparator; //使用委托
        //protected IComparer<E> comparator;

        public AVLTree() : this(null)
        {
        }
        public AVLTree(Comparison<E> comparer)
        {
            this.comparator = comparer;
        }

        private class AVLNode<E> : Node<E>
        {
            int height = 1;
            public AVLNode(E element, Node<E> parent) : base(element, parent) { }

            //获取平衡因子
            public int balanceFactor()
            {
                int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
                int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
                return leftHeight - rightHeight;//比较的是左右两边总高度的差值
            }
            public void updateHeight()
            {
                int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
                int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
                height = 1 + Math.Max(leftHeight, rightHeight);
            }
            /// <summary>
            /// 返回子树高度最大的节点
            /// </summary>
            /// <returns></returns>
            public Node<E> tallerChild()
            {
                int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
                int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
                if (leftHeight > rightHeight) return left;
                if (rightHeight > leftHeight) return right;
                return isLeftChild() ? left : right;
            }
        }
        protected override Node<E> createNode(E element, Node<E> parent)
        {
            return new AVLNode<E>(element, parent);
        }

        private void updateHeight(Node<E> node)
        {
            ((AVLNode<E>)node).updateHeight();
        }

        private bool isBalanced(Node<E> node)
        {
            return Math.Abs(((AVLNode<E>)node).balanceFactor()) <= 1;
        }

        protected override void afterAdd(Node<E> node)
        {
            //只有祖先节点会失衡
            while (node.parent != null)
            {
                node = node.parent;
                if (isBalanced(node))
                {
                    // 平衡则更新高度
                    updateHeight(node);
                }
                else
                {
                    // 恢复平衡
                    rebalance(node);
                    break;
                }
            }
        }
        protected override void afterRemove(Node<E> node)
        {
            while (node != null)
            {
                if (isBalanced(node))
                {
                    // 平衡则更新高度
                    updateHeight(node);
                }
                else
                {
                    // 恢复平衡    会影响到祖先，所以全部要平衡
                    rebalance(node);
                }
                node = node.parent;
            }
        }

        /// <summary>
        /// 恢复平衡  四种情况，加入位置在左边的右边， 就向右边和左边旋转 即RL
        /// </summary>
        /// <param name="node"></param>
        private void rebalance(Node<E> grad)
        {
            //取最高节点，就是不平衡的地方
            Node<E> parent = ((AVLNode<E>)grad).tallerChild();
            Node<E> node = ((AVLNode<E>)parent).tallerChild();
            if (parent.isLeftChild())
            {
                if (node.isLeftChild())
                {
                    // LL
                    rotateRight(grad);
                }
                else
                {
                    // LR
                    rotateLeft(parent);
                    rotateRight(grad);
                }
            }
            else
            {
                if (node.isLeftChild())
                {
                    // RL
                    rotateRight(parent);
                    rotateLeft(grad);
                }
                else
                {
                    // RR
                    rotateLeft(grad);
                }
            }
        }
        /// <summary>
        /// 旋转，只改变三个节点
        /// </summary>
        /// <param name="grand"></param>
        private void rotateLeft(Node<E> grand)
        {
            Node<E> parent = grand.right;
            Node<E> leftChild = parent.left;
            grand.right = leftChild;
            parent.left = grand;

            afterRotate(grand, parent, leftChild);
        }
        private void rotateRight(Node<E> grand)
        {
            Node<E> parent = grand.left;
            Node<E> rightChild = parent.right;
            grand.left = rightChild;
            parent.right = grand;

            afterRotate(grand, parent, rightChild);
        }



        private void afterRotate(Node<E> grand, Node<E> parent, Node<E> child)
        {
            parent.parent = grand.parent;

            // parent上位
            parent.parent = grand.parent;// 换爹
            // 爹也要换儿子
            if (grand.isLeftChild())
            {
                grand.parent.left = parent;
            }
            else if (grand.isRightChild())
            {
                grand.parent.right = parent;
            }
            else
            {
                //根节点
                root = parent;
            }

            //更新祖先节点
            grand.parent = parent;

            //更新child
            if (child != null)
            {
                child.parent = grand;
            }


            //更新高度
            updateHeight(grand);
            updateHeight(parent);
        }

    }
}
